from copy import deepcopy
import networkx as nx
import matplotlib.pyplot as plt
from networkx.drawing.nx_agraph import graphviz_layout


class CNOT_Node:
    # 节点
    #   pair: (q0, q1)
    #   uid:  节点id
    #   next_node: 下一个节点 [(CNOT_Node, distance距离),(CNOT_Node, distance距离)] 至多有两个
    #   pre_node: 上一个节点 [(CNOT_Node, distance距离),(CNOT_Node, distance距离)] 至多有两个
    def __init__(self, cid, qubit_pair: tuple, next_node: list, pre_node: list, pre_distance: list[int], next_distance: list[int]):
        self.cid = cid
        self.qubit_pair = qubit_pair
        self.next_node: list[CNOT_Node] = next_node
        self.pre_node: list[CNOT_Node] = pre_node
        self.pre_distance: list[int] = pre_distance
        self.next_distance: list[int] = next_distance

    def __deepcopy__(self, memo):
        '''创建新实例，不递归复制 pre_node 和 next_node'''
        new_instance = type(self).__new__(type(self))
        memo[id(self)] = new_instance  # 防止重复拷贝

        # 手动复制基本字段
        for k, v in self.__dict__.items():
            if k in ['pre_node', 'next_node']:
                continue  # 跳过节点引用，防止无限递归
            setattr(new_instance, k, deepcopy(v, memo))

        # 初始化为空列表，后续由 Circuit_DAG 处理连接
        new_instance.pre_node = []
        new_instance.next_node = []

        return new_instance


class Circuit_DAG:
    def __init__(self, qasm_path: str):
        self.qasm_path = qasm_path
        self.dag_nx: nx.DiGraph = None
        self.latest_cnot: dict[int, CNOT_Node] = {}
        self.total_cnot: dict[int, int] = {}
        self.cx_gates = []
        self.qubit_num = 0
        self.root_node = CNOT_Node(cid="root", qubit_pair=None, next_node=[], pre_node=[
        ], pre_distance=[0, 0], next_distance=[0, 0])
        self.parse_qasm_file()
        
    def __deepcopy__(self, memo):
        new_dag = type(self).__new__(type(self))
        memo[id(self)] = new_dag
        if 'dag_nx' in self.__dict__:
            new_dag.dag_nx = self.dag_nx.copy()  # 使用 nx 的 copy 方法
        if 'root_node' in self.__dict__:
            new_dag.root_node = deepcopy(self.root_node, memo)
        if 'cx_gates' in self.__dict__:
            new_dag.cx_gates = deepcopy(self.cx_gates, memo)
        if 'latest_cnot' in self.__dict__:
            new_dag.latest_cnot = {k: deepcopy(v, memo)
                                for k, v in self.latest_cnot.items()}
        if 'total_cnot' in self.__dict__:
            new_dag.total_cnot = deepcopy(self.total_cnot, memo)
        if 'qubit_num' in self.__dict__:
            new_dag.qubit_num = self.qubit_num
        if 'cnot_count' in self.__dict__:
            new_dag.cnot_count = self.cnot_count
        if 'qasm_path' in self.__dict__:
            new_dag.qasm_path = self.qasm_path
        return new_dag

    # 获取DAG的第一层
    def get_front_layer(self):
        front_nodes = [
            node for node in self.dag_nx.nodes() if self.dag_nx.in_degree(node) == 0]
        return front_nodes

    def parse_qasm_file(self):
        file_path = self.qasm_path
        with open(file_path, 'r') as file:
            lines = file.readlines()

        # qubit num
        qubit_line = next(line for line in lines if line.startswith('qubit['))
        num = int(qubit_line.split('[')[1].split(']')[0])
        self.qubit_num = num

        for i in range(num):
            self.latest_cnot[i] = self.root_node

        cnot_count = 0
        # all cnot gates
        for line in lines[3:]:
            if line.startswith('cx'):
                # extract the edge
                parts = line.split('q[')
                if len(parts) >= 3:
                    num1 = int(parts[1].split(']')[0])
                    num2 = int(parts[2].split(']')[0])
                    tmp_partial_cnot = (
                        num1, num2) if num1 < num2 else (num2, num1)
                    self.cx_gates.append(tmp_partial_cnot)
                    if (tmp_partial_cnot in self.total_cnot):
                        self.total_cnot[tmp_partial_cnot] += 1
                    else:
                        self.total_cnot[tmp_partial_cnot] = 1
                    dag_cnot_node = CNOT_Node(
                        cid=cnot_count, qubit_pair=(num1, num2), next_node=[], pre_node=[], pre_distance=[0, 0], next_distance=[0, 0])
                    cnot_count += 1
                    for edge_node_id in [num1, num2]:
                        dag_cnot_node.pre_node.append(
                            self.latest_cnot[edge_node_id])
                        self.latest_cnot[edge_node_id].next_node.append(
                            dag_cnot_node)
                        self.latest_cnot[edge_node_id] = dag_cnot_node
            else:
                if ('q[' in line):
                    parts = line.split('q[')
                    qubit_id = int(parts[1].split(']')[0])
                    if (self.latest_cnot[qubit_id].qubit_pair == None):
                        continue
                    if (self.latest_cnot[qubit_id].qubit_pair[0] == qubit_id):
                        self.latest_cnot[qubit_id].next_distance[0] += 1
                    else:
                        self.latest_cnot[qubit_id].next_distance[1] += 1
        self.cnot_count = cnot_count
        # sort self.total_cnot by key's value
        self.total_cnot = sorted(self.total_cnot.items())
        # self.visualize_dag()
        self.convert_to_nx()

    def convert_to_nx(self):
        G = nx.DiGraph()
        visited = set()
        queue: list[CNOT_Node] = [self.root_node]

        # DAG边：节点a --- 阻挡门数量、时间权重 ---> 节点b

        while queue:
            current = queue.pop(0)
            if current.cid in visited:
                continue
            visited.add(current.cid)
            G.add_node(current.cid, label=current.qubit_pair,
                       time_weight=0, process_done_time=0)
            for next_node in current.next_node:
                if (current.qubit_pair == None):
                    G.add_edge(current.cid, next_node.cid,
                               weight=0)
                    queue.append(next_node)
                    continue
                if (next_node.qubit_pair[0] == current.qubit_pair[0] or next_node.qubit_pair[1] == current.qubit_pair[0]):
                    G.add_edge(current.cid, next_node.cid,
                               weight=current.next_distance[0])
                else:
                    G.add_edge(current.cid, next_node.cid,
                               weight=current.next_distance[1])
                queue.append(next_node)
        self.dag_nx = G

    def visualize_dag(self):
        G = nx.DiGraph()
        visited = set()
        queue: list[CNOT_Node] = [self.root_node]

        while queue:
            current = queue.pop(0)
            if current.cid in visited:
                continue
            visited.add(current.cid)
            G.add_node(current.cid, label=current.qubit_pair)
            for next_node in current.next_node:
                if (current.qubit_pair == None):
                    G.add_edge(current.cid, next_node.cid, weight=0)
                    queue.append(next_node)
                    continue
                if (next_node.qubit_pair[0] == current.qubit_pair[0] or next_node.qubit_pair[1] == current.qubit_pair[0]):
                    G.add_edge(current.cid, next_node.cid,
                               weight=current.next_distance[0])
                else:
                    G.add_edge(current.cid, next_node.cid,
                               weight=current.next_distance[1])
                queue.append(next_node)
        pos = graphviz_layout(G, prog='dot', args='-Grankdir=LR')
        labels = nx.get_node_attributes(G, 'label')
        edge_labels = nx.get_edge_attributes(G, 'weight')
        plt.figure(figsize=(100, 80))
        # plt.figure(figsize=(10, 8))
        nx.draw(G, pos, with_labels=True, labels=labels, node_size=2000,
                node_color='skyblue', font_size=10, font_weight='bold')
        nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)
        plt.title("DAG Visualization")
        plt.show()
